'''
Created on 24 juin 2011

@author: EliX
'''
from multiprocessing import (Event, Lock)
from sage.combinat.backtrack import SearchForest
from sage.misc.abstract_method import abstract_method
import threading
import multiprocessing
import copy
import Queue

CONSTSYNC = 10000

class SearchForest2(SearchForest):
    
    def __init__(self, roots = None, 
                 children = None, 
                 post_process = None,
                 statistic = None, 
                 reduce_function = None, 
                 reduce_init = None, 
                 algorithm = 'depth', 
                 category=None) :
        SearchForest.__init__(self, roots, children, post_process, algorithm, category)
        if post_process is None :
            self.post_process = lambda x: x
        if statistic is not None :
            self.statistic = statistic
        if reduce_function is not None :
            self.reduce_function = reduce_function
        self._res = None
        self._init = copy.copy(reduce_init)

    @abstract_method
    def statistic(self, o):
        r''' TODO '''
    
    @abstract_method
    def reduce_function(self):
        r''' TODO '''
        
    def _initialisation(self):
        self._res = copy.copy(self._init)
        self._queue = Queue.Queue()
        n = max (multiprocessing.cpu_count(), 1)
        self._event = Event()
        self._varFIN = 0
        self._varTMP = 0
        self._lock = Lock()
        self._thds = [Travelator(self, 2**i) for i in range(n)]
        self._maskFIN = 0
        for i in range(n) :
            self._maskFIN = self._maskFIN | 2**i
    
    def _distrib(self):
        self._queue.put( self.roots() )
        for i in range( len(self._thds) ) :
            nodes = self._queue.get()
            for n in nodes :
                update(n, self, self)
                self._queue.put( iter( self.children(n) ) )
        
    def _start(self):
        for t in self._thds :
            t.start()
 
    def _assemblage(self):
        for t in self._thds :
            self._res = self.reduce_function(self._res, t._res)
    
    def run(self):
        # initialisation
        self._initialisation()
        # distribution
        self._distrib()
        # demarrage du parcours
        self._start()
        # attente de la fin du parcours
        self._event.wait()
        #print "EVENT"
        # assemblage des résultats
        self._assemblage()
        # retour de la fonction
        return self._res
    
class Travelator(threading.Thread):
    def __init__(self, obj, pos):
        threading.Thread.__init__(self)
        self.pos = pos
        self._res = copy.copy(obj._init)
        self._liste = []
        self._obj = obj
        self._run = True
        # print pos

    def run(self):
        i = 0
        obj = self._obj
        while self._run :
            node = None
            if len(self._liste) > 0 :
                # cas classique
                node = self._liste.pop()
            else :
                # cas liste vide
                try :
                    node = obj._queue.get(False)
                except :
                    # si la queue est vide
                    # declaration en dormeur
                    with obj._lock :
                        obj._varTMP = obj._varTMP | self.pos
                    # tant qu'il n'y a rien dans la queue
                    while node is None and (obj._varTMP & self.pos) is not 0 :
                        try : node = obj._queue.get(False)
                        except :
                            if obj._varTMP == obj._maskFIN :
                                #print "phase2 :: ", self.pos
                                with obj._lock :
                                    obj._varFIN = obj._varFIN | self.pos
                                while obj._varFIN != obj._maskFIN and (obj._varFIN & self.pos) is not 0 and obj._varTMP == obj._maskFIN :
                                    pass
                                if obj._varFIN == obj._maskFIN and obj._varTMP == obj._maskFIN :
                                    # print "FIN :: ", self.pos
                                    if self.pos == 1 :
                                        obj._event.set()
                                    self._run = False
                                    break
            
            if node is None :
                # print "None"
                continue
            
            if obj._varTMP != 0 :
                with obj._lock :
                    obj._varTMP = 0
                    obj._varFIN = 0
                    obj._event.clear()
            
            for n in node :
                update(n, obj, self)
                it = iter( obj.children(n) )
                i = i + 1
                if i == CONSTSYNC :
                    i = 0
                    obj._queue.put( it )
                else : self._liste.append( it )

def update(n, obj, th) :
    ''' S'occupe de toute la partie maj des valeurs ''reduce_function'' 
    '''
    node = obj.post_process(n)      
    if node is None :
        return

    # maj des valeurs
    th._res = obj.reduce_function(th._res, obj.statistic(node))
